//
// Created by haoy on 2017/7/12.
//

#ifndef RTC_XSWITCH2_NALU_HPP
#define RTC_XSWITCH2_NALU_HPP

#include <iostream>
#include <list>
#include <sstream>
#include <stdint.h>

#include "xcutil.h"

namespace h264 {
enum nal_type {
    reserved_0 = 0,
    slice_non_idr,
    slice_data_partition_a,
    slice_data_partition_b,
    slice_data_partition_c,
    slice_idr = 5,
    supplement_enhancement_information,
    SPS,
    PPS,
    access_unit_delimiter,
    end_of_sequence = 10,
    end_of_stream,
    filler_data,
    sequence_parameter_set_extension,
    prefix_nal_unit,
    subset_sequence_parameter_set = 15,
    depth_parameter_set,
    reserved_17,
    reserved_18,
    slice_layer_without_partitioning,
    slice_layer_extension = 20,
    slice_layer_extension_depth,
    reserved_22,
    reserved_23,
    STAP_A,
    single_time_aggregation_packet_b = 25,
    multi_time_aggregation_packet_16,
    multi_time_aggregation_packet_24,
    fragment_unit_a,
    fragment_unit_b,
    reserved_30,
    reserved_31,
};

extern std::string nal_type_name(nal_type t);

class nal_header {
public:
    int read(uint8_t *data, int size) {
        if (size < 1) {
            return 0;
        }
        f = data[0] >> 7;
        nri = data[0] >> 5 & 0x03;
        type = data[0] & 0x1F;
        return 1;
    }

    int write(uint8_t *data, int size) const {
        if (size < 1) {
            return 0;
        }
        data[0] = f << 7 | nri << 5 | type;
        return 1;
    }

    int size() const {
        return 1;
    }

    std::string to_string() const {
        std::ostringstream w;
        w << "{\"forbidden_zero_bit\":" << f << ",\"nal_ref_idc\":" << nri << ",\"nal_type\":\""
          << nal_type_name((nal_type)type) << "\"}";
        return w.str();
    }

    int f;
    int nri;
    int type;
};

class nal_single {
public:
    int read(uint8_t *data, int sz) {
        if (sz < 1) {
            return 0;
        }
        payload_size = sz;
        if (payload_size < 0) {
            return 0;
        }
        payload = data;
        return sz;
    }

    int write(uint8_t *data, int sz) const {
        if (sz < size()) {
            return 0;
        }
        memcpy(data, payload, payload_size);
        return payload_size;
    }

    int size() const {
        return payload_size;
    }

    std::string to_string() const {
        std::ostringstream w;
        w << "{\"payload_size\":" << payload_size << "}";
        return w.str();
    }

    uint8_t *payload;
    int payload_size;
};

class nal_stap_unit {
public:
    int read(uint8_t *data, int size) {
        if (size < 2) {
            return 0;
        }
        payload_size = be_get_u16(data) - 1;
        if (size < payload_size + 2 + 1) {
            return 0;
        }
        int header_read = header.read(data + 2, size - 2);
        if (header_read == 0) {
            return 0;
        }

        payload = data + 3;
        return payload_size + 3;
    }

    int write(uint8_t *data, int size) const {
        if (size < payload_size + 2 + header.size()) {
            return 0;
        }
        be_set_u16(payload_size, data);
        header.write(data + 2, size - 2);
        memcpy(data + 2 + header.size(), payload, payload_size);
        return payload_size + 2 + header.size();
    }

    int size() const {
        return payload_size + 2 + header.size();
    }

    std::string to_string() const {
        std::ostringstream w;
        w << "{\"stap_unit_header\":" << header.to_string() << ",";
        w << "\"payload_size\":" << payload_size << "}";
        return w.str();
    }

    int payload_size;
    nal_header header;
    uint8_t *payload;
};

class nal_stap_a {
public:
    int read(uint8_t *data, int s) {
        units.clear();
        int offset = 0;
        while (s - offset) {
            nal_stap_unit u;
            int unit_read = u.read(data + offset, s - offset);
            // could be padding
            if (unit_read == 0) {
                break;
            }
            offset += unit_read;
            units.emplace_back(u);
        }
        return offset;
    }

    int write(uint8_t *data, int s) const {
        if (s < size()) {
            return 0;
        }
        int offset = 0;
        for (auto &u : units) {
            int unit_write = u.write(data + offset, s - offset);
            if (unit_write == 0) {
                return 0;
            }
            offset += unit_write;
        }
        return offset;
    }

    int size() const {
        int s = 0;
        for (auto &u : units) {
            s += u.size();
        }
        return s;
    }

    std::string to_string() const {
        std::ostringstream w;
        w << "{";
        w << "\"units\":[";
        for (auto &u : units) {
            w << u.to_string() << ",";
        }
        w << "]}";
        return w.str();
    }

    std::list<nal_stap_unit> units;
};

class nal_fu_header {
public:
    int read(uint8_t *data, int sz) {
        if (sz < 1) {
            return 0;
        }
        s = data[0] >> 7;
        e = data[0] >> 6 & 0x01;
        r = data[0] >> 5 & 0x01;
        type = data[0] & 0x1F;
        return 1;
    }

    int write(uint8_t *data, int sz) const {
        if (sz < 1) {
            return 0;
        }
        data[0] = s << 7 | e << 6 | r << 5 | type;
        return 1;
    }

    int size() const {
        return 1;
    }

    std::string to_string() const {
        std::ostringstream w;
        w << "{\"start\":" << s << ",\"end\":" << e << ",\"reserved\":" << r << ",\"type\":\""
          << nal_type_name((nal_type)type) << "\"}";

        return w.str();
    }

    int s;
    int e;
    int r;
    int type;
};

class nal_fu_a {
public:
    int read(uint8_t *data, int sz) {
        int header_read = header.read(data, sz);
        if (header_read == 0) {
            return 0;
        }
        payload_size = sz - header_read;
        if (payload_size < 0) {
            return 0;
        }

        payload = data + header_read;
        return sz;
    }

    int write(uint8_t *data, int sz) const {
        if (sz < size()) {
            return 0;
        }
        int offset = 0;
        int header_write = header.write(data + offset, sz - offset);
        if (header_write == 0) {
            return 0;
        }
        offset += header_write;
        memcpy(data + offset, payload, payload_size);
        offset += payload_size;
        return offset;
    }

    int size() const {
        return header.size() + payload_size;
    }

    std::string to_string() const {
        std::ostringstream w;
        w << "{\"fu_header\":" << header.to_string() << ",\"payload_size\"=" << payload_size << "}";
        return w.str();
    }

    nal_fu_header header;
    uint8_t *payload;
    int payload_size;
};


class nal_packet {
public:
    int read(uint8_t *data, int sz) {
        int header_read = header.read(data, sz);
        if (header_read == 0) {
            return 0;
        }
        int payload_read = 0;
        switch (header.type) {
            case nal_type::STAP_A:
                payload_read = stap_a.read(data + header_read, sz - header_read);
                break;
            case nal_type::fragment_unit_a:
                payload_read = fu_a.read(data + header_read, sz - header_read);
                break;
            default:
                payload_read = single.read(data + header_read, sz - header_read);
        }
        if (payload_read == 0) {
            return 0;
        }
        return header_read + payload_read;
    }

    int write(uint8_t *data, int sz) const {
        if (sz < size()) {
            return 0;
        }
        int header_write = header.write(data, sz);
        if (header_write == 0) {
            return 0;
        }
        int payload_write = 0;
        switch (header.type) {
            case nal_type::STAP_A:
                payload_write = stap_a.write(data + header_write, sz - header_write);
                break;
            case nal_type::fragment_unit_a:
                payload_write = fu_a.write(data + header_write, sz - header_write);
                break;
            default:
                payload_write = single.write(data + header_write, sz - header_write);
        }
        if (payload_write == 0) {
            return 0;
        }
        return header_write + payload_write;
    }

    int size() const {
        switch (header.type) {
            case nal_type::STAP_A:
                return header.size() + stap_a.size();
            case nal_type::fragment_unit_a:
                return header.size() + fu_a.size();
        }
        return header.size() + single.size();
    }

    std::string to_string() const {
        std::ostringstream w;
        w << "{";
        w << "\"indicator\":" << header.to_string() << ",";
        w << "\"payload\":";
        switch (header.type) {
            case nal_type::STAP_A:
                w << stap_a.to_string();
                break;
            case nal_type::fragment_unit_a:
                w << fu_a.to_string();
                break;
            default:
                w << single.to_string();
        }
        w << "}";
        return w.str();
    }

    nal_header header;
    nal_single single;
    nal_stap_a stap_a;
    nal_fu_a fu_a;
};

} // namespace h264
#endif //RTC_XSWITCH2_NALU_HPP
